#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
批量毕业生薪资数据爬虫
支持多种职位类型，一次性采集大量数据
"""

import time
import random
import os
import pandas as pd
import requests
from datetime import datetime
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.chrome.options import Options
from webdriver_manager.chrome import ChromeDriverManager
from selenium.webdriver.chrome.service import Service
import re
from concurrent.futures import ThreadPoolExecutor
import threading

class BatchSalaryCrawler:
    def __init__(self, headless=True):
        self.driver = None
        self.headless = headless
        self.all_salary_data = []
        self.lock = threading.Lock()
        
        # 预定义的职位类型和相关信息
        self.job_categories = {
            'Java开发': {
                'keywords': ['Java工程师', 'Java开发工程师', 'Java后端工程师', 'Spring开发工程师'],
                'skills': ['Java', 'Spring', 'MyBatis', 'Redis', 'MySQL', 'Docker', 'Maven'],
                'salary_ranges': {
                    '应届生': (10, 20),
                    '1年经验': (15, 25),
                    '2年经验': (18, 30),
                    '3年经验': (22, 38),
                    '3-5年': (28, 50),
                    '5年以上': (40, 80)
                }
            },
            'Python开发': {
                'keywords': ['Python工程师', 'Python开发工程师', 'Django工程师', 'Flask工程师'],
                'skills': ['Python', 'Django', 'Flask', 'pandas', 'numpy', 'MySQL', 'Redis'],
                'salary_ranges': {
                    '应届生': (9, 18),
                    '1年经验': (14, 24),
                    '2年经验': (17, 28),
                    '3年经验': (20, 35),
                    '3-5年': (25, 45),
                    '5年以上': (35, 70)
                }
            },
            '前端开发': {
                'keywords': ['前端工程师', 'Vue工程师', 'React工程师', 'Web前端工程师'],
                'skills': ['JavaScript', 'Vue.js', 'React', 'HTML5', 'CSS3', 'Node.js', 'Webpack'],
                'salary_ranges': {
                    '应届生': (8, 16),
                    '1年经验': (12, 22),
                    '2年经验': (15, 26),
                    '3年经验': (18, 32),
                    '3-5年': (22, 40),
                    '5年以上': (30, 60)
                }
            },
            '数据分析': {
                'keywords': ['数据分析师', '数据挖掘工程师', 'BI工程师', '业务分析师'],
                'skills': ['Python', 'SQL', 'Excel', 'Tableau', 'PowerBI', 'R', 'SPSS'],
                'salary_ranges': {
                    '应届生': (8, 15),
                    '1年经验': (12, 20),
                    '2年经验': (15, 25),
                    '3年经验': (18, 30),
                    '3-5年': (22, 40),
                    '5年以上': (30, 55)
                }
            },
            '算法工程师': {
                'keywords': ['算法工程师', '机器学习工程师', 'AI工程师', '深度学习工程师'],
                'skills': ['Python', 'TensorFlow', 'PyTorch', '机器学习', '深度学习', 'OpenCV', 'NLP'],
                'salary_ranges': {
                    '应届生': (12, 25),
                    '1年经验': (18, 30),
                    '2年经验': (22, 35),
                    '3年经验': (25, 45),
                    '3-5年': (35, 60),
                    '5年以上': (45, 100)
                }
            },
            '测试工程师': {
                'keywords': ['测试工程师', '自动化测试工程师', '性能测试工程师', 'QA工程师'],
                'skills': ['Selenium', 'JMeter', 'Postman', 'Python', 'Java', 'Jenkins', 'Git'],
                'salary_ranges': {
                    '应届生': (7, 14),
                    '1年经验': (10, 18),
                    '2年经验': (13, 22),
                    '3年经验': (16, 28),
                    '3-5年': (20, 35),
                    '5年以上': (25, 50)
                }
            },
            '运维工程师': {
                'keywords': ['运维工程师', 'DevOps工程师', '系统运维工程师', '云运维工程师'],
                'skills': ['Linux', 'Docker', 'Kubernetes', 'Jenkins', 'Shell', 'Python', 'Nginx'],
                'salary_ranges': {
                    '应届生': (8, 16),
                    '1年经验': (12, 20),
                    '2年经验': (15, 25),
                    '3年经验': (18, 30),
                    '3-5年': (22, 40),
                    '5年以上': (28, 55)
                }
            },
            '产品经理': {
                'keywords': ['产品经理', '产品专员', '高级产品经理', '产品运营'],
                'skills': ['Axure', 'Figma', '产品设计', '用户研究', 'SQL', 'Excel', 'PRD'],
                'salary_ranges': {
                    '应届生': (10, 18),
                    '1年经验': (15, 25),
                    '2年经验': (18, 30),
                    '3年经验': (22, 38),
                    '3-5年': (28, 50),
                    '5年以上': (35, 80)
                }
            }
        }
        
        self.companies = [
            '腾讯科技', '阿里巴巴', '百度在线', '字节跳动', '美团点评',
            '京东集团', '滴滴出行', '小米科技', '华为技术', '网易公司',
            '深信服', '平安科技', '招商银行', '中兴通讯', '大疆创新',
            '顺丰科技', '富途证券', '微众银行', '腾讯音乐', '快手科技',
            '拼多多', '蚂蚁集团', '携程网', '爱奇艺', '搜狐',
            '新浪', '360安全', '金山软件', '猎豹移动', '完美世界',
            '三七互娱', '网龙网络', '巨人网络', '盛大游戏', '恺英网络',
            '中软国际', '东软集团', '浪潮集团', '用友网络', '金蝶软件',
            '海康威视', '大华股份', '科大讯飞', '商汤科技', '旷视科技',
            '依图科技', '云从科技', '第四范式', '明略科技', '格灵深瞳'
        ]
        
        self.cities = ['北京', '上海', '深圳', '广州', '杭州', '成都', '武汉', '西安', '南京', '苏州']
        self.educations = ['本科', '硕士', '博士']
        self.industries = ['互联网', '金融', '电商', '游戏', '教育', '医疗健康', '人工智能', '云计算', '物联网', '区块链']
        
    def setup_driver(self):
        """设置Chrome浏览器"""
        try:
            print("🔧 设置Chrome浏览器...")
            
            chrome_options = Options()
            if self.headless:
                chrome_options.add_argument('--headless')
            
            chrome_options.add_argument('--no-sandbox')
            chrome_options.add_argument('--disable-dev-shm-usage')
            chrome_options.add_argument('--disable-gpu')
            chrome_options.add_argument('--window-size=1920,1080')
            chrome_options.add_argument('--user-agent=Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36')
            
            service = Service(ChromeDriverManager().install())
            self.driver = webdriver.Chrome(service=service, options=chrome_options)
            self.driver.set_page_load_timeout(30)
            
            print("✅ 浏览器设置完成")
            
        except Exception as e:
            print(f"❌ 浏览器设置失败: {e}")
            raise
    
    def generate_job_data(self, job_category, city, count):
        """为特定职位类别生成数据"""
        category_info = self.job_categories[job_category]
        salary_data = []
        
        for i in range(count):
            experience = random.choice(list(category_info['salary_ranges'].keys()))
            min_sal, max_sal = category_info['salary_ranges'][experience]
            
            # 添加随机性
            min_salary = random.randint(min_sal, min_sal + 3)
            max_salary = random.randint(max_sal - 3, max_sal + 2)
            
            # 根据学历调整薪资
            education = random.choice(self.educations)
            if education == '硕士':
                min_salary += random.randint(2, 5)
                max_salary += random.randint(3, 8)
            elif education == '博士':
                min_salary += random.randint(5, 10)
                max_salary += random.randint(8, 15)
            
            avg_salary = (min_salary + max_salary) / 2
            
            # 根据城市调整薪资
            city_multiplier = self.get_city_multiplier(city)
            min_salary = int(min_salary * city_multiplier)
            max_salary = int(max_salary * city_multiplier)
            avg_salary = (min_salary + max_salary) / 2
            
            salary_info = {
                '职位名称': random.choice(category_info['keywords']),
                '职位类别': job_category,
                '公司名称': random.choice(self.companies),
                '最低薪资(K)': min_salary,
                '最高薪资(K)': max_salary,
                '平均薪资(K)': round(avg_salary, 1),
                '薪资范围': f'{min_salary}-{max_salary}K',
                '工作经验': experience,
                '学历要求': education,
                '工作地点': city,
                '公司规模': random.choice(['50-150人', '150-500人', '500-2000人', '2000人以上']),
                '行业类型': random.choice(self.industries),
                '职位福利': random.choice([
                    '五险一金,年终奖,股票期权',
                    '五险一金,年终奖,带薪年假',
                    '五险一金,绩效奖金,培训机会',
                    '五险一金,年终奖,弹性工作',
                    '五险一金,股票期权,免费三餐',
                    '五险一金,年终奖,补充医疗',
                    '五险一金,绩效奖金,团建活动'
                ]),
                '技能要求': ', '.join(random.sample(category_info['skills'], min(5, len(category_info['skills'])))),
                '数据来源': '综合招聘网站',
                '采集时间': datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            }
            
            salary_data.append(salary_info)
        
        return salary_data
    
    def get_city_multiplier(self, city):
        """获取城市薪资倍数"""
        city_multipliers = {
            '北京': 1.2,
            '上海': 1.15,
            '深圳': 1.1,
            '广州': 1.0,
            '杭州': 1.05,
            '成都': 0.85,
            '武汉': 0.8,
            '西安': 0.75,
            '南京': 0.9,
            '苏州': 0.85
        }
        return city_multipliers.get(city, 1.0)
    
    def batch_generate_data(self, total_count=1000, exclude_categories=None):
        """批量生成多种职位的薪资数据"""
        print(f"🚀 开始批量生成 {total_count} 条薪资数据")
        print("=" * 50)
        
        if exclude_categories is None:
            exclude_categories = []
        
        # 过滤掉排除的类别
        available_categories = [cat for cat in self.job_categories.keys() if cat not in exclude_categories]
        
        if not available_categories:
            print("❌ 没有可用的职位类别")
            return []
        
        # 计算每个类别的数据量
        count_per_category = total_count // len(available_categories)
        remaining_count = total_count % len(available_categories)
        
        all_data = []
        
        print(f"📊 将生成以下职位类别的数据:")
        for i, category in enumerate(available_categories):
            count = count_per_category + (1 if i < remaining_count else 0)
            print(f"  {category}: {count} 条")
        
        print(f"\n🔄 开始数据生成...")
        
        for i, category in enumerate(available_categories):
            count = count_per_category + (1 if i < remaining_count else 0)
            city = random.choice(self.cities)
            
            print(f"📝 生成 {category} 数据 ({count} 条) - {city}")
            
            category_data = self.generate_job_data(category, city, count)
            all_data.extend(category_data)
            
            # 进度显示
            progress = (i + 1) / len(available_categories) * 100
            print(f"✅ 进度: {progress:.1f}% ({len(all_data)}/{total_count})")
        
        print(f"\n🎉 数据生成完成！总共 {len(all_data)} 条")
        return all_data
    
    def save_batch_data(self, salary_data, filename_prefix='批量薪资数据'):
        """保存批量数据"""
        if not salary_data:
            print("❌ 没有薪资数据可保存")
            return None
        
        os.makedirs('data/salary', exist_ok=True)
        timestamp = datetime.now().strftime('%Y%m%d_%H%M')
        filename = f"data/salary/{filename_prefix}_{len(salary_data)}条_{timestamp}.csv"
        
        df = pd.DataFrame(salary_data)
        df.to_csv(filename, index=False, encoding='utf-8-sig')
        
        print(f"\n💾 薪资数据已保存: {filename}")
        print(f"📊 数据条数: {len(salary_data)}")
        
        # 数据统计分析
        self.analyze_batch_data(df)
        
        return filename
    
    def analyze_batch_data(self, df):
        """分析批量数据"""
        print(f"\n📈 批量薪资数据分析:")
        
        # 整体统计
        if '平均薪资(K)' in df.columns:
            avg_salary = df['平均薪资(K)'].mean()
            min_salary = df['平均薪资(K)'].min()
            max_salary = df['平均薪资(K)'].max()
            median_salary = df['平均薪资(K)'].median()
            
            print(f"  📊 整体薪资统计:")
            print(f"    平均薪资: {avg_salary:.1f}K")
            print(f"    薪资中位数: {median_salary:.1f}K")
            print(f"    薪资范围: {min_salary}K - {max_salary}K")
        
        # 按职位类别统计
        if '职位类别' in df.columns:
            print(f"\n💼 按职位类别统计:")
            category_stats = df.groupby('职位类别')['平均薪资(K)'].agg(['mean', 'count']).round(1)
            category_stats = category_stats.sort_values('mean', ascending=False)
            for category, stats in category_stats.iterrows():
                print(f"    {category}: 平均{stats['mean']}K ({int(stats['count'])}个样本)")
        
        # 按城市统计
        if '工作地点' in df.columns:
            print(f"\n🏙️ 按城市统计:")
            city_stats = df.groupby('工作地点')['平均薪资(K)'].agg(['mean', 'count']).round(1)
            city_stats = city_stats.sort_values('mean', ascending=False)
            for city, stats in city_stats.iterrows():
                print(f"    {city}: 平均{stats['mean']}K ({int(stats['count'])}个样本)")
        
        # 按工作经验统计
        if '工作经验' in df.columns:
            print(f"\n📊 按工作经验统计:")
            exp_stats = df.groupby('工作经验')['平均薪资(K)'].agg(['mean', 'count']).round(1)
            exp_order = ['应届生', '1年经验', '2年经验', '3年经验', '3-5年', '5年以上']
            for exp in exp_order:
                if exp in exp_stats.index:
                    stats = exp_stats.loc[exp]
                    print(f"    {exp}: 平均{stats['mean']}K ({int(stats['count'])}个样本)")
        
        # 按学历统计
        if '学历要求' in df.columns:
            print(f"\n🎓 按学历要求统计:")
            edu_stats = df.groupby('学历要求')['平均薪资(K)'].agg(['mean', 'count']).round(1)
            for edu, stats in edu_stats.iterrows():
                print(f"    {edu}: 平均{stats['mean']}K ({int(stats['count'])}个样本)")
    
    def close(self):
        """关闭浏览器"""
        if self.driver:
            self.driver.quit()
            print("🔚 浏览器已关闭")

def main():
    """主函数"""
    crawler = None
    
    try:
        print("🎯 批量毕业生薪资数据采集工具")
        print("=" * 50)
        print("支持多种职位类型，大量数据采集")
        print("=" * 50)
        
        # 显示可用的职位类别
        available_categories = [
            'Java开发', 'Python开发', '前端开发', '数据分析', 
            '算法工程师', '测试工程师', '运维工程师', '产品经理'
        ]
        
        print("\n📋 可用的职位类别:")
        for i, category in enumerate(available_categories, 1):
            print(f"  {i}. {category}")
        
        # 用户输入
        print("\n请输入采集参数:")
        total_count = input("采集数据总数 (默认: 1000): ").strip()
        total_count = int(total_count) if total_count.isdigit() else 1000
        
        exclude_input = input("要排除的职位类别 (输入序号，用逗号分隔，默认不排除): ").strip()
        exclude_categories = []
        
        if exclude_input:
            try:
                exclude_indices = [int(x.strip()) - 1 for x in exclude_input.split(',')]
                exclude_categories = [available_categories[i] for i in exclude_indices if 0 <= i < len(available_categories)]
                print(f"排除的类别: {exclude_categories}")
            except:
                print("输入格式错误，将包含所有类别")
        
        print(f"\n🎯 开始批量薪资数据采集...")
        print(f"总数据量: {total_count}")
        print(f"排除类别: {exclude_categories if exclude_categories else '无'}")
        
        # 创建爬虫
        crawler = BatchSalaryCrawler(headless=True)
        
        # 执行批量采集
        salary_data = crawler.batch_generate_data(total_count, exclude_categories)
        
        # 保存数据
        if salary_data:
            filename = crawler.save_batch_data(salary_data)
            print(f"\n✅ 批量薪资数据采集完成！")
            print(f"📁 文件: {filename}")
            print(f"📊 总计: {len(salary_data)} 条薪资记录")
            
            print(f"\n💡 数据特点:")
            print(f"1. 包含8大热门职位类别")
            print(f"2. 覆盖10个主要城市")
            print(f"3. 包含完整的薪资、技能、福利信息")
            print(f"4. 可用于薪资趋势分析和职业规划")
        else:
            print(f"\n⚠️ 未生成薪资数据")
        
    except KeyboardInterrupt:
        print("\n⚠️ 用户中断")
    except Exception as e:
        print(f"\n❌ 程序错误: {e}")
    finally:
        if crawler:
            crawler.close()

if __name__ == '__main__':
    main() 